/*******************************************************************************
    Copyright (c) 2011 Yahoo! Inc. All rights reserved.

    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License. See accompanying LICENSE file.

    The Initial Developer of the Original Code is Shravan Narayanamurthy.
******************************************************************************/
/*
 * Controller.cpp
 *
 * The Controller script for learntopics
 *
 * Creates the relevant Model_Builder as per the
 * given model & mode.
 *
 * Uses the director to build the model using the
 * above created Model_Builder
 *
 * Asks the builder to create the output after the
 * model is built
 *
 *  Created on: 06-Jan-2011
 *      
 */

#include "Main_flags_define.h"
#include "glog/logging.h"
#include "TopicLearner/Unigram_Model_Training_Builder.h"
#include "TopicLearner/Unigram_Model_Synchronized_Training_Builder.h"
#include "TopicLearner/Unigram_Model_Testing_Builder.h"
#include "TopicLearner/Unigram_Model_Streaming_Builder.h"
#include "Model_Director.h"
#include "Model.h"
#include "DocumentReader.h"
#include "document.pb.h"

using namespace std;

int main(int argc, char *argv[]) {
    google::InitGoogleLogging(argv[0]);
    google::InstallFailureSignalHandler();
    string
            usage(
                    "\n\nThis program learns a topic model from a corpus which has been formatted using \'formatter\'.\n\n"
                        "Inputs are the files generated by \'formatter\'. So you need to run this in the directory that contains these files.\n\n"
                        "You can customize where the inputs are taken from and where outputs go to using the \'inputprefix\' flag.\n\n"
                        "You can configure many aspects of the topic model using the many optional flags. The most important ones are:\n\n"
                        "The model that has to be used specified by \'model\' flag\n"
                        "The number of topics you want to learn specified by \'topics\' flag\n"
                        "The number of iterations the program should run controlled by \'iter\' flag\n"
                        "The mode the program runs in - test or teststream or training - controlled by \'test\' and \'teststream\' flag\n\n"
                        "You can also run this as a part of the multi-machine set-up.\n"
                        "Take a look at docs/html/index.html for more info.\n\n");
    string
            cmd_usage(
                    (string) argv[0]
                            + " [--inputprefix=<inputprefix>]\n"
                                "         [--iter=<#iterations_to_run> --topics=<#topics_to_learn>]\n"
                                "         [-test -teststream --dumpprefix=<dump_of_Model>]\n"
                                "         [--alpha=<topic_conjugate> --beta=<word_conjugate>]\n"
                                "         [--burnin=<burnin> --printloglikelihood=<logLikelihood_interval> --optimizestats=<param_optimize_interval>]\n"
                                "         [--servers=<memcached_servers>] ---chkptdir=<dir_to_check_point> --chkptinterval=<#check_point_interval>\n");

    usage += cmd_usage;
    google::SetUsageMessage(usage);

    for (int i = 0; i < argc; i++) {
        std::string arg = argv[i];
        if (strcmp(argv[i], "--help") == 0) {
            google::ShowUsageWithFlagsRestrict(argv[0], "Main");
            exit(0);
        }
    }

    google::ParseCommandLineFlags(&argc, &argv, true);

    google::SetCommandLineOptionWithMode("minloglevel", "1",
            google::SET_FLAG_IF_DEFAULT);
    google::SetCommandLineOptionWithMode("stderrthreshold", "1",
            google::SET_FLAG_IF_DEFAULT);

    const char* pwd = google::StringFromEnv("PWD", "/tmp");
    google::SetCommandLineOptionWithMode("log_dir", pwd,
            google::SET_FLAG_IF_DEFAULT);

    LOG(WARNING)
            << "----------------------------------------------------------------------";
    LOG(WARNING) << "Log files are being stored at " << pwd
            << "/learnTopics.*";
    LOG(WARNING)
            << "----------------------------------------------------------------------";

    string flagsInp = google::CommandlineFlagsIntoString();

    LOG(INFO) << flagsInp << endl;

    //Do parameter checking here
    if (!FLAGS_test && !FLAGS_teststream) {
        //Training mode
        if (google::GetCommandLineFlagInfoOrDie("servers").is_default) {
            //Single machine mode
            LOG(WARNING) << "You have chosen single machine training mode";
        } else {
            //Multi machine mode
            LOG(WARNING) << "You have chosen multi machine training mode";
        }
    } else {
        //Test mode
        if (google::GetCommandLineFlagInfoOrDie("servers").is_default) {
            //Single machine mode
            LOG(WARNING) << "You have chosen single machine testing mode";
        } else {
            //Multi machine mode
            LOG(WARNING) << "You have chosen multi machine testing mode";
        }
    }

    Model_Builder* builder;
    if (FLAGS_model == Model::UNIGRAM) {
        if (!FLAGS_test && !FLAGS_teststream) {
            if (!google::GetCommandLineFlagInfoOrDie("servers").is_default)
                builder = new Unigram_Model_Synchronized_Training_Builder();
            else
                builder = new Unigram_Model_Training_Builder();
        } else {
            if (FLAGS_test)
                builder = new Unigram_Model_Testing_Builder();
            else
                builder = new Unigram_Model_Streaming_Builder();
        }
    }

    Model_Director director;
    Model& model = director.build_model(*builder);
    if (!FLAGS_test && !FLAGS_teststream)
        LOG(WARNING) << "Model has been learnt";
    if (!FLAGS_teststream)
        builder->create_output();

    delete builder;
}
